import os
import logging
from typing import Optional
from dotenv import load_dotenv

# 加载环境变量
load_dotenv()

logger = logging.getLogger(__name__)

class LLMService:
    """LLM服务类，支持本地离线模型和API调用"""
    
    def __init__(self):
        self.model = None
        self.tokenizer = None
        self.is_initialized = False
        self.mode = os.getenv("LLM_MODE", "local")  # local 或 api
        self.knowledge_base = None  # 初始化知识库
        
        # 本地模型配置
        self.local_model_name = os.getenv("LOCAL_MODEL_NAME", "offline")
        
        # API配置（备选）
        self.api_key = os.getenv("LLM_API_KEY")
        self.api_base = os.getenv("LLM_API_BASE", "https://api.openai.com/v1")
        self.api_model_name = os.getenv("LLM_MODEL_NAME", "gpt-3.5-turbo")
    
    async def initialize(self):
        """初始化LLM服务"""
        try:
            logger.info(f"正在初始化LLM服务，模式: {self.mode}")
            
            if self.mode == "local":
                await self._initialize_local_model()
            else:
                await self._initialize_api_mode()
            
            self.is_initialized = True
            logger.info("LLM服务初始化完成")
            
        except Exception as e:
            logger.error(f"LLM服务初始化失败: {e}")
            # API模式作为备选
            await self._initialize_api_mode()
            self.is_initialized = True
    
    async def _initialize_local_model(self):
        """初始化本地离线模型"""
        try:
            logger.info("正在初始化离线本地模型...")
            
            # 构建离线知识库
            self.knowledge_base = {
                "transformer": {
                    "title": "Attention Is All You Need",
                    "authors": ["Vaswani, A.", "Shazeer, N.", "Parmar, N.", "Uszkoreit, J.", "Jones, L.", "Gomez, A. N.", "Kaiser, L.", "Polosukhin, I."],
                    "year": 2017,
                    "abstract": "提出了一种基于自注意力机制的神经网络架构，完全替代了循环和卷积神经网络。",
                    "key_points": [
                        "自注意力机制允许模型关注输入序列中的不同位置",
                        "编码器-解码器架构支持并行计算",
                        "位置编码保持序列顺序信息"
                    ],
                    "citations": 80000,
                    "arxiv_id": "1706.03762"
                },
                "bert": {
                    "title": "BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding",
                    "authors": ["Devlin, J.", "Chang, M. W.", "Lee, K.", "Toutanova, K."],
                    "year": 2018,
                    "abstract": "提出了一种新的语言表示模型BERT，通过预训练深度双向Transformer在多个NLP任务上取得突破。",
                    "key_points": [
                        "掩码语言模型实现双向上下文理解",
                        "预训练+微调的统一框架",
                        "在11个NLP任务上刷新记录"
                    ],
                    "citations": 50000,
                    "arxiv_id": "1810.04805"
                },
                "gpt": {
                    "title": "Improving Language Understanding by Generative Pre-Training",
                    "authors": ["Radford, A.", "Narasimhan, K.", "Salimans, T.", "Sutskever, I."],
                    "year": 2018,
                    "abstract": "展示了生成式预训练在自然语言理解任务上的有效性。",
                    "key_points": [
                        "无监督预训练+有监督微调",
                        "基于Transformer解码器的架构",
                        "在多个任务上优于专门模型"
                    ],
                    "citations": 30000,
                    "arxiv_id": None
                },
                "resnet": {
                    "title": "Deep Residual Learning for Image Recognition",
                    "authors": ["He, K.", "Zhang, X.", "Ren, S.", "Sun, J."],
                    "year": 2015,
                    "abstract": "提出残差学习框架解决深度网络训练难题。",
                    "key_points": [
                        "残差连接解决梯度消失问题",
                        "支持训练极深网络",
                        "在ImageNet等数据集上取得突破"
                    ],
                    "citations": 120000,
                    "arxiv_id": "1512.03385"
                },
                "machine learning": {
                    "title": "Machine Learning",
                    "authors": ["Mitchell, T. M."],
                    "year": 1997,
                    "abstract": "机器学习的经典教材，系统介绍了机器学习的基本概念和方法。",
                    "key_points": [
                        "监督学习、无监督学习、强化学习",
                        "特征工程和模型评估",
                        "过拟合与欠拟合问题"
                    ],
                    "citations": 50000,
                    "arxiv_id": None
                }
            }
            
            logger.info("✓ 离线本地模型初始化完成")
            
        except Exception as e:
            logger.error(f"离线模型初始化失败: {e}")
            await self._initialize_api_mode()
    
    async def _initialize_api_mode(self):
        """初始化API模式"""
        logger.info("使用API模式")
        self.mode = "api"
        
        if not self.api_key:
            logger.warning("未设置LLM_API_KEY，将使用演示模式")
    
    async def generate_response(self, prompt: str, max_tokens: int = 1000) -> str:
        """生成响应"""
        try:
            if self.mode == "local" and hasattr(self, 'knowledge_base'):
                return await self._generate_local_response(prompt, max_tokens)
            else:
                return await self._generate_api_response(prompt, max_tokens)
                
        except Exception as e:
            logger.error(f"生成响应失败: {e}")
            return self._generate_fallback_response(prompt)
    
    async def _generate_local_response(self, prompt: str, max_tokens: int) -> str:
        """使用离线本地模型生成响应"""
        try:
            prompt_lower = prompt.lower()
            
            # 关键词匹配
            matched_topics = []
            for topic, info in self.knowledge_base.items():
                if topic in prompt_lower:
                    matched_topics.append((topic, info))
            
            if matched_topics:
                response = "基于我的离线知识库，我找到了以下相关信息：\n\n"
                for topic, info in matched_topics:
                    response += f"**{info['title']}** ({info['year']})\n"
                    response += f"作者: {', '.join(info['authors'][:3])}等\n"
                    response += f"摘要: {info['abstract']}\n\n"
                    response += "关键贡献:\n"
                    for point in info['key_points']:
                        response += f"- {point}\n"
                    if info['arxiv_id']:
                        response += f"\nArxiv链接: https://arxiv.org/abs/{info['arxiv_id']}\n"
                    response += f"引用次数: {info['citations']:,}\n\n"
                
                response += "如需更详细信息，建议在Arxiv上搜索相关关键词。"
                return response
            else:
                return f"""基于您的问题，我分析了相关学术领域：

**问题分析**: {prompt}

**建议研究方向**:
1. 在Arxiv上搜索相关关键词获取最新论文
2. 查看该领域的重要会议和期刊
3. 关注相关研究团队的最新工作

*当前为离线模式，基于预构建知识库提供回答。*"""
            
        except Exception as e:
            logger.error(f"离线模型生成失败: {e}")
            return await self._generate_api_response(prompt, max_tokens)
    
    async def _generate_api_response(self, prompt: str, max_tokens: int) -> str:
        """使用API生成响应"""
        try:
            if not self.api_key:
                return self._generate_demo_response(prompt)
            
            # 这里可以集成OpenAI API或其他API服务
            # 目前使用演示响应
            return self._generate_demo_response(prompt)
            
        except Exception as e:
            logger.error(f"API调用失败: {e}")
            return self._generate_fallback_response(prompt)
    
    def _generate_demo_response(self, prompt: str) -> str:
        """生成演示响应"""
        prompt_lower = prompt.lower()
        
        if any(word in prompt_lower for word in ['transformer', 'attention', '注意力']):
            return """Transformer是一种基于自注意力机制的神经网络架构，由Vaswani等人在2017年的论文《Attention Is All You Need》中提出。

主要特点：
1. **自注意力机制**：允许模型直接关注输入序列中的任意位置
2. **并行计算**：相比RNN可以完全并行处理序列
3. **位置编码**：通过正弦编码保持序列位置信息

Transformer已成为现代NLP的基础架构，催生了BERT、GPT等重要模型。"""
        
        elif any(word in prompt_lower for word in ['bert', 'gpt', 'llm']):
            return """大语言模型（LLM）的发展历程：

**BERT（2018）**：
- 双向编码器架构
- 掩码语言模型预训练
- 在多个NLP任务上取得突破

**GPT系列**：
- 自回归生成模型
- 从GPT到GPT-4参数规模持续增长
- 展现了强大的few-shot和zero-shot学习能力

当前LLM技术在文本理解、代码生成、对话系统等领域都有广泛应用。"""
        
        else:
            return f"""基于您的问题，我检索了相关学术文献并进行分析：

**分析要点**：
- 该问题涉及多个研究领域的交叉
- 相关论文提供了不同的理论视角
- 需要结合具体应用场景进行深入探讨

**建议**：
如需更详细的回答，请提供更具体的问题描述或相关背景信息。

*当前为演示模式，实际部署后将基于真实LLM模型提供更准确的回答。*"""
    
    def _generate_fallback_response(self, prompt: str) -> str:
        """生成备用响应"""
        return "抱歉，当前无法生成回答。请检查LLM服务配置或稍后重试。"
    
    def is_ready(self) -> bool:
        """检查服务是否就绪"""
        return self.is_initialized
    
    async def cleanup(self):
        """清理资源"""
        pass